package com.example.ssrbac.config;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.context.annotation.Configuration;
import org.springframework.messaging.simp.config.MessageBrokerRegistry;
import org.springframework.web.socket.config.annotation.EnableWebSocketMessageBroker;
import org.springframework.web.socket.config.annotation.StompEndpointRegistry;
import org.springframework.web.socket.config.annotation.WebSocketMessageBrokerConfigurer;
import org.springframework.web.socket.server.support.DefaultHandshakeHandler;
import org.springframework.web.socket.server.HandshakeInterceptor;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.http.server.ServletServerHttpRequest;
import javax.servlet.http.HttpServletRequest;
import java.util.Map;

@Configuration
@EnableWebSocketMessageBroker
public class WebSocketConfig implements WebSocketMessageBrokerConfigurer {
    private static final Logger logger = LoggerFactory.getLogger(WebSocketConfig.class);

    @Override
    public void configureMessageBroker(MessageBrokerRegistry config) {
        logger.info("配置WebSocket消息代理");
        // 配置消息代理
        // 设置消息代理的前缀
        config.enableSimpleBroker("/topic");
        // 设置客户端发送消息的前缀
        config.setApplicationDestinationPrefixes("/app");
        logger.info("WebSocket消息代理配置完成");
    }

    @Override
    public void registerStompEndpoints(StompEndpointRegistry registry) {
        logger.info("注册WebSocket端点");
        // 注册STOMP端点
        registry.addEndpoint("/ws")
                .setAllowedOriginPatterns("*")  // 允许所有来源
                .addInterceptors(new HandshakeInterceptor() {
                    @Override
                    public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response,
                            WebSocketHandler wsHandler, Map<String, Object> attributes) throws Exception {
                        logger.debug("WebSocket握手开始: {}", request.getURI());
                        if (request instanceof ServletServerHttpRequest) {
                            HttpServletRequest servletRequest = ((ServletServerHttpRequest) request).getServletRequest();
                            String token = servletRequest.getParameter("token");
                            logger.debug("WebSocket握手token: {}", token);
                            if (token == null || token.isEmpty()) {
                                logger.warn("WebSocket握手失败: token为空");
                                return false;
                            }
                            attributes.put("token", token);
                            logger.debug("WebSocket握手成功");
                        }
                        return true;
                    }

                    @Override
                    public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response,
                            WebSocketHandler wsHandler, Exception exception) {
                        if (exception != null) {
                            logger.error("WebSocket握手后发生错误", exception);
                        } else {
                            logger.debug("WebSocket握手完成");
                        }
                    }
                })
                .withSockJS();  // 启用SockJS回退选项
        logger.info("WebSocket端点注册完成");
    }
} 